from sklearn.datasets import load_iris
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split
import sklearn as sk
import torch as pt
from torch.utils.data import TensorDataset, DataLoader
import numpy as np
import matplotlib.pyplot as plt
import sys

np.random.seed(777)
pt.manual_seed(777)

# 1.	使用sklearn配合pytorch完成以下操作（每题10分）
# (1)	数据处理
# ①	读取iris数据集
x, y = load_iris(return_X_y=True)
scaler = StandardScaler()
x = scaler.fit_transform(x)
print('x', x.shape)
print('y', y.shape)

# ②	自定义超参数数值
ALPHA = 0.001
N_EPOCHS = 500
BATCH_SIZE = 32
M, N = x.shape
N_CLS = len(np.unique(y))

# ③	将数据7:3切分
x_train, x_test, y_train, y_test = train_test_split(x, y, train_size=0.7, random_state=777)
M_TRAIN, _ = x_train.shape
x_train = pt.from_numpy(x_train).float()
y_train = pt.from_numpy(y_train).long()
x_test = pt.from_numpy(x_test).float()
y_test = pt.from_numpy(y_test).long()

# ④	放入dataloder
ds_train = TensorDataset(x_train, y_train)
dl_train = DataLoader(ds_train, batch_size=BATCH_SIZE, shuffle=True)

# ⑤	创建全连接层，两个隐藏层，神经元数量分别是10,5
# ⑥	每个隐藏层后用relu激活，dropout设定0.2
L1 = 100  # 为了效果，添加神经元数量
L2 = 50  # 为了效果，添加神经元数量
model = pt.nn.Sequential(
    pt.nn.Linear(N, L1),
    pt.nn.ReLU(),
    pt.nn.Dropout(0.2),
    pt.nn.Linear(L1, L2),
    pt.nn.ReLU(),
    pt.nn.Dropout(0.2),
    pt.nn.Linear(L2, N_CLS)
)

# ⑦	使用交叉熵作为代价函数
criterion = pt.nn.CrossEntropyLoss()
optim = pt.optim.Adam(params=model.parameters(), lr=ALPHA)


def acc(h, y):
    return h.argmax(1).eq(y.long()).float().mean()


# ⑧	训练模型，每100批次打印损失值和准确率
cost_history = np.zeros(N_EPOCHS)
GROUP = 100
for epoch in range(N_EPOCHS):
    cost_avg = 0.
    acc_avg = 0.
    for i, (bx, by) in enumerate(dl_train):
        model.train(True)
        optim.zero_grad()
        h = model(bx)
        cost = criterion(h, by)
        cost.backward()
        optim.step()
        model.train(False)
        cost = cost.detach().numpy()
        accv = acc(h, by).detach().numpy()
        cost_avg += cost
        acc_avg += accv
    cost_avg /= i + 1
    acc_avg /= i + 1
    cost_history[epoch] = cost
    if epoch % GROUP == 0:
        print(f'epoch#{epoch + 1}: cost = {cost}, acc = {accv}')
if epoch % GROUP != 0:
    print(f'epoch#{epoch + 1}: cost = {cost}, acc = {accv}')

# ⑨	绘制损失变化曲线
plt.plot(cost_history)
plt.title('Cost value in each epoch')

# ⑩	打印最终准确率
h = model(x_test)
accv = acc(h, y_test).detach().numpy()
print(f'测试集最终准确率: {accv}')

# Finally show all plotting
plt.show()
